# -*- coding: utf-8 -*-
import datetime

from sqlalchemy import func
from sqlalchemy.orm import sessionmaker, Query

from py_engine import engine
import logging

from py_table import User

logger = logging.getLogger(__name__)


class SessionContext(object):
    """
    数据库连接会话上下文
    """

    def __init__(self, _engine):
        self.engine = _engine

    def __enter__(self):
        _Session = sessionmaker(bind=self.engine, autoflush=False)
        self.session = _Session()
        return self.session

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.session.close()


def test_session_context():
    with SessionContext(engine) as session:
        query: Query = session.query(User)
        user: User = query.get(1)
        user.user_birthday = datetime.date(2023, 11, 11)
        session.commit()


def test_session():
    with SessionContext(engine) as session:
        # 聚合函数
        max_user_id = session.query(func.max(User.user_id)).one_or_none()[0]
        max_user_id = max_user_id if max_user_id else 0  # 注意表空时返回None

        # 插入
        user = User(user_name=f'cch{max_user_id + 1}')
        session.add(user)
        # 批量插入
        user_list = [User(user_name=f'chench{max_user_id + 1}'), User(user_name=f'chchen{max_user_id + 1}'), ]
        session.add_all(user_list)

        # 查询
        query: Query = session.query(User)
        # 查询所有记录
        all_user = query.all()
        # 根据主键查询
        get_user = query.get(max_user_id)

        # 更新
        query_first = query.first()
        print('update:', query_first.user_id)
        query_first.user_name = 'cch' + str(max_user_id)

        # 删除
        query_first = query.first()
        print('delete:', query_first.user_id)
        session.delete(query_first)

        try:
            session.prepare()
        except Exception as e:
            logger.error(e)
            session.rollback()
        else:
            session.commit()


def test_scoped_session():
    # sessionmaker返回的Session不是线程安全的，可以使用scoped_session处理
    session: Session = scoped_session(sessionmaker(engine))  # noqa


if __name__ == '__main__':
    test_session_context()
    test_session()
